import torch
import torch.nn as nn

def Hermitian(X): 
    """ torch矩阵共轭转置 """
    X = torch.real(X) - 1j * torch.imag(X)
    return X.transpose(-1, -2)

def eta(R, sigma_t2, K, G, M, theta):  
    """ 
        R: [batch, G, K] 
        G: 原始稀疏数据长度
        K: 子载波数也就是MMV的数量
        M: 为观测维度 
    """
    theta1 = torch.abs(theta[0])
    theta2 = theta[1]
    pi_t = 1 + sigma_t2 / theta1
    psi_t = 0.5 * K * torch.log(torch.tensor(1 + theta1 / sigma_t2).cuda()) + theta2
    
    aa = (psi_t.reshape(R.shape[0], 1, 1) - torch.sum(torch.abs(R * R), [2]).reshape(R.shape[0], R.shape[1], 1) / ((2 * sigma_t2 * pi_t).reshape(R.shape[0], 1, 1)))
    frac_div = torch.sigmoid(-aa) / pi_t.reshape(R.shape[0], 1, 1)

    H_res = R * frac_div
    b = torch.sum(frac_div, [1, 2]) / M
    return H_res, b


class LAMP_Net(nn.Module): # 单层
    def __init__(self, Theta, Nc, N, pilots): 
        #H [batch,G,K]为角度域信道, Theta [Nc,L,N]为感知矩阵 N为原始稀疏数据长度，T为迭代次数，Nc为子载波数也就是MMV的数量，L为观测维度
        super(LAMP_Net,self).__init__()
        Theta_real = torch.tensor(Theta["real"]).cuda()
        Theta_imag = torch.tensor(Theta["imag"]).cuda()
        Theta = torch.complex(Theta_real, Theta_imag)
        self.B = torch.nn.Parameter(Hermitian(Theta))    #[Nc,N,L]
        self.pilots = pilots
        self.N = N
        self.Nc = Nc
        self.theta = torch.nn.Parameter(torch.tensor([1.0,2.0]))
        
    def forward(self, Y, Theta): 
        # 输入: Y: [batch, Nc, L], 要改成 [batch, L, Nc]
        # N = ((torch.randn(H.shape[0], M, K).cuda() + 1j * torch.randn(H.shape[0], M, K).cuda()) * (sqrt(sigma2/2))).cuda()
        # Y = torch.matmul(self.A, H) + N
        Y = Y.permute(0,2,1).reshape(Y.shape[0], self.pilots, self.Nc)

        H_res = torch.zeros(Y.shape[0], self.N, self.Nc).cuda() # [batch, N, Nc]
        V = Y # [batch, L, Nc]
        
        B = self.B
        theta = self.theta

        BV = B @ V.permute(0,2,1).reshape(Y.shape[0], self.Nc, self.pilots, 1) # [ batch, Nc ,N, 1]
        BV = BV.permute(0,2,1,3).reshape(Y.shape[0], self.N, self.Nc)          # [batch, N, Nc]

        R = H_res + BV  # [batch, N, Nc]
        sigma_t2 = torch.sum(torch.abs(V*V),[1,2]) / (self.pilots * self.Nc)
        H_res,b = eta(R, sigma_t2, self.Nc, self.N, self.pilots, theta)

        HH = Theta @ H_res.permute(0,2,1).reshape(Y.shape[0], self.Nc, self.N, 1) # [batch, Nc, L, 1]
        HH = HH.permute(0,2,1,3).reshape(Y.shape[0], self.pilots, self.Nc)        # [batch, L, Nc]

        V = Y - HH + b.reshape(-1,1,1) * V  # # [batch, L, Nc]
        
        return H_res, V


class LAMP_MuNet(nn.Module): # 多层
    def __init__(self, Theta, Nc, N, pilots, net): 
        # H: [batch,G,K]为角度域信道, Theta: [Nc,L,N]为感知矩阵, N: 原始稀疏数据长度, T: 迭代次数, Nc: 子载波数也就是MMV的数量, L: 观测维度
        super(LAMP_MuNet,self).__init__()
        if isinstance(net, str): # 路径
            self.net = torch.load(net)
        else:
            self.net = net
        for p in self.net.parameters():
            (p.requires_grad) = False
        
        self.pilots = pilots
        self.N = N
        self.Nc = Nc
        Theta_real = torch.tensor(Theta["real"]).cuda()
        Theta_imag = torch.tensor(Theta["imag"]).cuda()
        Theta = torch.complex(Theta_real, Theta_imag)
        self.B = torch.nn.Parameter(Hermitian(Theta))    # [Nc, N, L]
        self.theta = torch.nn.Parameter(torch.tensor([1.0,2.0]))
        
        
    def forward(self, Y, Theta): 
        # 输入: Y [batch, Nc, L], 要改成 [batch, L, Nc]
        H_res, V = self.net(Y, Theta) 

        Y = Y.permute(0,2,1).reshape(Y.shape[0], self.pilots, self.Nc)

        B = self.B
        theta = self.theta

        BV = B @ V.permute(0,2,1).reshape(Y.shape[0], self.Nc, self.pilots, 1) # [batch, Nc, N, 1]
        BV = BV.permute(0,2,1,3).reshape(Y.shape[0], self.N, self.Nc) # [batch, N, Nc]

        R = H_res + BV
        sigma_t2 = torch.sum(torch.abs(V*V),[1,2]) / (self.pilots * self.Nc)
        H_res, b = eta(R, sigma_t2, self.Nc, self.N, self.pilots, theta)

        HH = Theta @ H_res.permute(0,2,1).reshape(Y.shape[0], self.Nc, self.N, 1)   # [batch, Nc, L, 1]
        HH = HH.permute(0,2,1,3).reshape(Y.shape[0], self.pilots, self.Nc) # [batch, L, Nc]

        V = Y - HH + b.reshape(-1,1,1) * V
        
        return H_res, V